import collections

import time
import datetime
import functools
import logging
import sys

from absl import app, flags

import pandas

import coin.strategy.mm.feed as feed
import coin.strategy.mm.tool.plot_util as putil
from coin.base.book.types import TradeSide, create_market_trade_event
from coin.base.datetime_util import to_timestamp_int
from coin.proto.coin_order_gateway_pb2 import OrderEvent, OrderGatewayLog
from coin.strategy.mm.tool.util import KuonaOrderManager, get_order_sign

import coin.strategy.mm.tool.archive_base as abase

LevelInfo = collections.namedtuple('LevelInfo', ['price', 'qty'])

FLAGS = flags.FLAGS


def plot_table(ax, df):
  cellDict = ax.table(cellText=df.values, colLabels=df.columns, cellLoc='center',
                      loc='center').get_celld()
  for i in range(len(df.index) + 1):
    cellDict[(i, 0)].set_width(0.06)
    for j in range(1, len(df.columns)):
      cellDict[(i, j)].set_width(0.05)


def iter_book(iteritems_book, n):
  if FLAGS.use_feed_cache:
    for i, level in enumerate(iteritems_book):
      if i >= n:
        break
      yield level
  else:
    for i, (_, level) in enumerate(iteritems_book):
      if i >= n:
        break
      yield level


BookShow = collections.namedtuple('BookShow', ['bidp', 'bidq', 'askp', 'askq'])


def book_to_show(book):
  nbook = 15
  if FLAGS.use_feed_cache:
    bidpqs = [(a.price, a.qty) for a in iter_book(book.iter_bid(), nbook)]
    askpqs = [(a.price, a.qty) for a in iter_book(book.iter_ask(), nbook)]
  else:
    bidpqs = [(a.price, a.qty) for a in iter_book(book.iteritems_bid(), nbook)]
    askpqs = [(a.price, a.qty) for a in iter_book(book.iteritems_ask(), nbook)]
  bidp, bidq = zip(*bidpqs)
  askp, askq = zip(*askpqs)
  return BookShow(bidp, bidq, askp, askq)


def show_to_df(show):
  nbook = 15

  df = pandas.DataFrame(
      collections.OrderedDict([('bidp', [None] * 15 + list(show.bidp)),
                               ('bidq', [None] * 15 + list(show.bidq)),
                               ('askp', list(show.askp[::-1]) + [None] * 15),
                               ('askq', list(show.askq[::-1]) + [None] * 15)]))
  return df


def update_table(table, df):
  for i in range(df.values.shape[0]):
    for j in range(df.values.shape[1]):
      text = df.values[i, j]
      table.get_celld()[(i + 1, j)]._text.set_text(text)


class MatplotlibEngine(object):
  def __init__(self):
    import matplotlib
    matplotlib.rcParams.update({'font.size': 7})
    import matplotlib.pyplot as plt
    self.plt = plt
    self.fig = None
    self.die = False
    self.table = None

  def press(self, event):
    if event.key == 'q':
      self.die = True

  def draw_df(self, df):
    if self.fig is None:
      self.fig = self.plt.figure(figsize=(6, 6))
      self.fig.canvas.mpl_connect('key_press_event', self.press)
      self.table = self.plt.table(cellText=df.values, colLabels=df.columns, loc='center')
      self.plt.axis("off")
    else:
      update_table(self.table, df)
      self.fig.canvas.draw()
    self.plt.waitforbuttonpress(timeout=-1)

    if self.die:
      self.plt.close()
      sys.exit(1)


class Scroll(object):
  def __init__(self):
    self.cursor = 0
    self.events = []

  def add_event(self, timestamp, show, trade, askqd, bidqd):
    if FLAGS.use_feed_cache and trade is not None:
      tradefeed = create_market_trade_event(timestamp,
                                            trade.symbol,
                                            trade.price,
                                            trade.qty,
                                            trade.side,
                                            None)
      self.events.append((timestamp, show, tradefeed, askqd, bidqd))
    else:
      self.events.append((timestamp, show, trade, askqd, bidqd))

  def move_impl(self, cursor, step):
    assert step in [1, -1]
    while 0 <= cursor + step < len(self.events):
      cursor += step
      cursor = max(0, cursor)
      cursor = min(len(self.events), cursor)
      if self.events[cursor][1] is not None:
        return 1, cursor
    return 0, cursor

  def move(self, step):
    ret, cursor = self.move_impl(self.cursor, step)
    if ret == 1:
      self.cursor = cursor
    return ret

  def get_bbo_qtydiff(self):
    return (self.events[self.cursor][3], self.events[self.cursor][4])

  def get_book(self):
    if self.events[self.cursor][1] is None:
      self.move(1)
    assert self.events[self.cursor][1] is not None, self.cursor
    return show_to_df(self.events[self.cursor][1])

  def get_trades(self):
    trades = []
    for event in self.events[self.cursor::-1]:
      if event[2] is not None:
        trades.append(event[2])
      if len(trades) >= 30:
        break
    return trades

  def get_timestamp(self):
    return self.events[self.cursor][0]


from coin.strategy.mm.tool.util import get_order_sign, KuonaOrderManager


class OrderSnapshot(object):
  def __init__(self):
    self.order_tpqs = {}
    self.fills = []
    self.position = 0

  def update_order_tpqs(self, order_tpqs):
    self.order_tpqs = dict(order_tpqs).copy()

  def update_fill(self, atuple):
    self.fills.append(atuple)
    self.fills = self.fills[-8:]

  def update_position(self, position):
    self.position = position

  def copy(self):
    return (self.order_tpqs.copy(), self.fills.copy(), self.position)


class OrderScroll(object):
  def __init__(self):
    self.tss = []
    self.snapshots = []

  def empty(self):
    return len(self.tss) == 1

  def update_snapshot(self, timestamp, snapshot):
    snapshot_new = snapshot.copy()
    if len(self.snapshots) > 0 and snapshot_new == self.snapshots[-1]:
      return
    self.tss.append(timestamp)
    self.snapshots.append(snapshot.copy())

  def get_snapshot(self, timestamp):
    import bisect
    idx = bisect.bisect_left(self.tss, timestamp)
    return self.snapshots[idx - 1]


class Replay(object):
  def __init__(self, product, engine, start_ts, end_ts):
    self.product = product
    # self.engine = MatplotlibEngine()
    self.engine = engine
    self.scroll = Scroll()
    self.ordersnapshot = OrderSnapshot()
    self.oscroll = OrderScroll()
    self.oscroll.update_snapshot(0, self.ordersnapshot)
    self.order_manager = KuonaOrderManager(FLAGS.exchange, product)
    self.prev_bid0 = None
    self.prev_ask0 = None
    self._start_ts = to_timestamp_int(start_ts or 0)
    self._end_ts = to_timestamp_int(end_ts or 0)

  def on_book_reset(self, book_builder_name, book_builder):
    for product in [self.product]:
      print(book_builder_name, book_builder)
      book_builder.subscribe(product, functools.partial(self.on_book, product))
      book_builder.subscribe_trade(product, functools.partial(self.on_trade, product))

  def on_book(self, product, book):
    if book.timestamp < self._start_ts:
      return
    if self._end_ts and self._end_ts < book.timestamp:
      return

    if product != self.product:
      return
    show = book_to_show(book)

    if self.engine.die:
      sys.exit(1)

    askqd = bidqd = 0

    if (self.prev_bid0 is not None and self.prev_bid0.price == book.bid0().price):
      bidqd = book.bid0().qty - self.prev_bid0.qty

    if (self.prev_ask0 is not None and self.prev_ask0.price == book.ask0().price):
      askqd = book.ask0().qty - self.prev_ask0.qty

    self.prev_bid0 = LevelInfo(price=book.bid0().price, qty=book.bid0().qty)
    self.prev_ask0 = LevelInfo(price=book.ask0().price, qty=book.ask0().qty)

    self.scroll.add_event(book.timestamp, show, None, askqd, bidqd)

    while self.engine.find_qty is not None:
      step_actual = self.scroll.move(self.engine.step)
      if step_actual == 0:
        return
      if (self.scroll.get_bbo_qtydiff()[0] == self.engine.find_qty
          or self.scroll.get_bbo_qtydiff()[1] == self.engine.find_qty):
        self.engine.find_qty = None
        break

    while self.engine.skip == 0 and self.engine.find_order == False:
      if self.engine.die:
        sys.exit(1)
      time.sleep(0.0001)

    while self.engine.skip != 0:
      step_actual = self.scroll.move(self.engine.step)
      self.engine.skip -= step_actual
      if step_actual == 0:
        break
    # print(self.scroll.cursor, self.engine.skip)

    while not self.oscroll.empty() and self.engine.find_order == True:
      step_actual = self.scroll.move(self.engine.step)
      if step_actual == 0:
        break
      order_tpqs, fills, self_pos = self.oscroll.get_snapshot(self.scroll.get_timestamp())
      if len(order_tpqs) > 0:
        self.engine.find_order = False
        break

    if self.engine.skip == 0 and self.engine.find_order == False:
      if not self.oscroll.empty():
        order_tpqs, fills, self_pos = self.oscroll.get_snapshot(self.scroll.get_timestamp())
      else:
        order_tpqs, fills, self_pos = [], [], 0

      trades = self.scroll.get_trades()
      self.engine.set_title(
          # str(pandas.DatetimeIndex([self.scroll.get_timestamp()])[0])
          datetime.datetime.fromtimestamp(self.scroll.get_timestamp() * 1e-9
                                         ).strftime("%Y%m%d %H:%M:%S.%f"))
      self.engine.update_table(self.scroll.get_book(),
                               trades,
                               order_tpqs,
                               self_pos,
                               fills,
                               self.product.symbol)
      time.sleep(0.01)

  def on_trade(self, product, trade):
    if trade.timestamp < self._start_ts:
      return
    if self._end_ts and self._end_ts < trade.timestamp:
      return
    if product != self.product:
      return
    self.scroll.add_event(trade.timestamp, None, trade, None, None)

  def on_log(self, timestamp, log):
    pb = OrderGatewayLog()
    try:
      pb.ParseFromString(log)
    except Exception:
      print(log)
      return
    if pb.type == OrderGatewayLog.ORDER_EVENT:
      event = pb.event
      if FLAGS.exchange in ['Bitflyer', 'Bitmex']:
        pass
      else:
        if event.internal_order_id == '':
          return
      symbol = event.symbol
      if (symbol != self.product and symbol != self.product.native_symbol
          and symbol != self.product.symbol):
        return
      timestamp = event.event_time
      price = event.order_price
      qty = event.order_qty
      side = event.order_side
      event_type = event.type
      fill_price = event.fill_price
      fill_qty = event.fill_qty
      sign = get_order_sign(event.order_side)

      if fill_qty > 0:
        self.ordersnapshot.update_fill((timestamp, fill_price, fill_qty * sign))

      if self.order_manager.on_event(event):
        self.ordersnapshot.update_order_tpqs(self.order_manager.order_tpqs)
    elif pb.type == OrderGatewayLog.POSITION:
      for each_position in pb.position.each_position:
        if each_position.symbol == self.product.symbol:
          self.ordersnapshot.update_position(each_position.net_position)
    else:
      return

    self.oscroll.update_snapshot(timestamp, self.ordersnapshot)


def launch_replay(engine, baseflags, fsflags, osflags, subsymbol):
  products = abase.subsymbols_to_products([subsymbol], baseflags)

  start_time, end_time = abase.get_time_range(baseflags)

  strategy = Replay(products[0], engine, start_ts=start_time, end_ts=end_time)

  print("order loading")
  if osflags.HasField('orderlog_machine'):
    abase.run_from_order_archive_base(baseflags, osflags, on_log=strategy.on_log)
  if flags.FLAGS.need_order:
    assert not strategy.oscroll.empty()

  print("order loading done")

  print("feed loading")
  abase.run_from_archive_base(baseflags=baseflags,
                              feed_sub_flags_proto=fsflags,
                              products=products,
                              on_book_reset=strategy.on_book_reset)


def launch(baseflags, fsflags, osflags):
  # launch_replay(None, baseflags, fsflags, osflags, symbol)
  # return

  from threading import Thread

  from PyQt5.QtWidgets import QApplication
  from coin.strategy.mm.tool.replay_pyqt_engine import PyQtEngine, set_palette
  app = QApplication([])
  engine = PyQtEngine(FLAGS.find_qty)

  subsymbols = None if flags.FLAGS.subsymbols is None else flags.FLAGS.subsymbols.split(",")

  thread_replay = Thread(target=launch_replay,
                         args=(
                             engine,
                             baseflags,
                             fsflags,
                             osflags,
                             subsymbols[0],
                         ))
  thread_replay.start()
  set_palette(app)
  app.exec_()
  return 0


def main(_):
  subsymbols = None if flags.FLAGS.subsymbols is None else flags.FLAGS.subsymbols.split(",")

  baseflags = abase.get_base_flags()
  products = abase.subsymbols_to_products(subsymbols, baseflags)
  baseflags.exchange, baseflags.market_type = subsymbols[0].split(":")[:2]
  fsflags = abase.get_feed_subscription_flags()
  osflags = abase.get_order_subscription_flags()

  launch(baseflags, fsflags, osflags)


def init_flags():
  abase.define_base_flags()
  abase.define_feed_archive_flags()
  abase.define_order_archive_flags()

  flags.FLAGS.feed_machine = 'feed-01.ap-northeast-1.aws'
  flags.DEFINE_string('subsymbols', None, 'subscript symbol name')
  flags.DEFINE_float('large_qty', None, '')
  flags.DEFINE_float('find_qty', None, '')
  flags.DEFINE_bool('need_order', False, '')

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')


if __name__ == '__main__':
  init_flags()
  app.run(main)
